# -*- coding: utf-8 -*-
# created on 2016/8/4

from mathsolver.functions.base import *
from sympy import Add, Mul, Pow, latex, S, log, sqrt, E, root
from mathsolver.functions.sympy_utils import get_all_child

POW2 = [4, 9, 16, 25, 36, 49, 86, 81, 100]


class BasicJisuan(BaseFunction):
    """
    BasicJisuan().solver(BasePoly("(((-2)))**(2002)+(((-2)))**(2001)"))
    BasicJisuan().solver(BasePoly("(-0.125)**5*4**10"))
    BasicJisuan().solver(BasePoly("log(14)-2*log((7)/(3))+log(7)-log(18)"))
    BasicJisuan().solver(BasePoly("(I)/((sqrt(3)+3*I))"))
    """
    @staticmethod
    def simplify_log(expr):
        flag = False
        args = []
        for arg in expr.args:
            if isinstance(arg, log):
                v = arg.args[0]
                _base = arg.args[1] if len(arg.args) > 1 else E
                nums = [n for n in POW2 if (v / n).is_Integer and v / n != 1]

                if nums:
                    n = nums[-1]
                    ratio = v / n
                    args.append(2 * log(sqrt(n), _base) + log(ratio, _base))
                    flag = True
                elif sqrt(v).is_Integer:
                    args.append(2 * log(sqrt(v), _base))
                    flag = True

                elif root(v, 3).is_Integer:
                    args.append(3 * log(root(v, 3), _base))
                    flag = True

            else:
                args.append(arg)

        if expr.is_Add:
            return flag, Add._from_args(args)
        elif expr.is_Mul:
            return flag, Mul._from_args(args)
        else:
            return False, expr

    @staticmethod
    def is_simple(arg, flag=True):
        if arg.is_Number:
            return True

        if "_" in str(arg):
            return False
        if str(arg).find(str('I')) >= 0:
            return False

        if get_all_child(arg, lambda xx: xx.is_Pow and xx.args[0] == -1):
            return False

        if get_all_child(arg, lambda xx: xx.is_Pow and xx.args[1] > 5):
            return False

        simple = arg.expand().simplify()
        if str(arg) == str(simple) and flag:
            return True

        if arg.is_Add and set(arg.args) == set(simple.args) and flag:
            return True

        """
        if arg.is_rational:
            if arg.is_Mul and simple.q == 1/arg.args[1] and simple.p == arg.args[0]:
                return True
            elif arg.is_Pow and simple.q == arg.args[0]:
                return True
        """
        return False

    def latex_step(self, res):
        if hasattr(res, "is_Add") and res.is_Add:
            symbol1 = " + "
            symbol2 = " "
            _step = self.latex_step(res.args[0])
            for r in res.args[1:]:
                if r.is_Number and r.is_negative:
                    _step += symbol2 + self.latex_step(r)
                else:
                    _step += symbol1 + self.latex_step(r)

        elif hasattr(res, "is_Mul") and res.is_Mul:
            _step = ""
            is_frac = False
            p_r = res.args[0]
            for r in res.args[1:]:
                if r.is_Pow and r.args[1] == -1:
                    temp = "\\frac{%s}{%s}" % (self.latex_step(p_r), self.latex_step(r.args[0]))
                    is_frac = True
                else:
                    temp = self.latex_step(p_r)
                    if p_r.is_Add:
                        temp = "\\left(%s\\right)" % temp
                    is_frac = False

                if _step:
                    _step += " \\times " + temp
                else:
                    _step += temp

                p_r = r

            if not is_frac:
                if r.is_Add:
                    r = "\\left(%s\\right)" % r
                _step += " * " + self.latex_step(r)
        elif hasattr(res, "is_Pow") and res.is_Pow:
            if res.args[1] == 1 / 2.:
                _step = "\\sqrt{%s}" % new_latex(res.args[0])
            else:
                _step = "{%s}^{%s}" % (new_latex(res.args[0]), new_latex(res.args[1]))
        else:
            _step = new_latex(res)

        return _step

    @staticmethod
    def func_compute(expr):
        return expr

    def calculate(self, expr):
        expr = self.func_compute(expr)
        if expr.is_Add:
            args = list(expr.args)
            first = args[0]
            second = args[1]
            if isinstance(first, log) and isinstance(second, log):
                _base = first.args[1] if len(first.args) > 1 else None
                if _base:
                    res = log(first.args[0] * second.args[0], _base).simplify()
                else:
                    res = log(first.args[0] * second.args[0]).simplify()
                add = Add._from_args([res] + args[2:])

            elif self.is_simple(first) and self.is_simple(second):
                res = Add(first, second).together()
                add = Add._from_args([res] + args[2:])

            for i, arg in enumerate(args):
                if not self.is_simple(arg):
                    add = Add._from_args(args[:i] + [self.calculate(arg)] + args[i + 1:])
                    break

            if str(add) == str(expr):
                return add.evalf()
            else:
                return add

        elif expr.is_Mul:
            args = list(expr.args)
            first = args[0]
            second = args[1]
            if first.is_Pow and second.is_Pow:
                min_pow = min(first.args[1], second.args[1])
                res = Pow(Mul(first.args[0], second.args[0]).simplify(), min_pow) * \
                      Pow(first.args[0], first.args[1] - min_pow) * Pow(second.args[0], second.args[1] - min_pow)
                mul = Mul._from_args([res] + args[2:])
            elif self.is_simple(first) and second.is_Pow and self.is_simple(second.args[0]) and second.args[1] == -1:
                res = Mul(first, second).simplify()
                mul = Mul._from_args([res] + args[2:])
            elif self.is_simple(first) and self.is_simple(second):
                res = Mul(first, second).expand().simplify()
                mul = Mul._from_args([res] + args[2:])
            else:
                for i, arg in enumerate(args):
                    if not self.is_simple(arg):
                        mul = Mul._from_args(args[:i] + [self.calculate(arg)] + args[i + 1:])
                        break
            if str(mul) == str(expr):
                return mul.evalf()
            else:
                return mul
        elif expr.is_Pow:
            if expr.args[0] == -1:
                if expr.args[1].is_even or (1 / expr.args[1]).is_even:
                    return S.One
                else:
                    return -S.One

            # 1/sqrt(3) - sqrt(2)
            if expr.is_Pow and expr.args[1] == -1:
                under = expr.args[0]
                if under.is_Add:
                    value = (under * (under.args[0] - under.args[1])).expand()
                    if value.is_Number:
                        self.label.add("含根式分数求逆")
                        return under / value
            if expr.args[1] == 0:
                return S.One
            elif not self.is_simple(expr.args[0]):
                return Pow(self.calculate(expr.args[0]), expr.args[1])
            elif expr.args[1] < 5:
                return expr.expand().simplify()
            else:
                assert False, "pow degree more than 4"
        else:
            return expr

    def solver(self, *args):
        assert not isinstance(args[0], BaseSinPoly), "sinpoly unacceptable"
        expr = args[0].sympify()
        num_step = 0
        assert len(expr.free_symbols) == 0, "no free symbols allowed"
        self.label.add("计算")

        self.steps.append(["", self.latex_step(expr)])
        while not self.is_simple(expr):
            expr = self.calculate(expr)
            self.steps.append(["", "=" + self.latex_step(expr)])
            num_step += 1
            if num_step < 20:
                break  # "too many steps"
                # print expr
        self.output.append(BaseNumber(expr))
        return self


if __name__ == '__main__':
    pass
